from typing import Dict, Any, List, Optional
from .base_agent import BaseAgent, AgentResult
import logging
from datetime import datetime

class AgentOrchestrator:
    """智能体协调器"""
    
    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}
        self.workflow_registry: Dict[str, List[str]] = {}
        self.logger = logging.getLogger(__name__)
        
        # 配置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
    
    def register_agent(self, agent: BaseAgent) -> None:
        """注册智能体"""
        agent_name = agent.get_name()
        if agent_name in self.agents:
            self.logger.warning(f"智能体 {agent_name} 已存在，将被覆盖")
        
        self.agents[agent_name] = agent
        self.logger.info(f"注册智能体: {agent_name} v{agent.get_version()}")
    
    def register_workflow(self, name: str, agent_sequence: List[str]) -> None:
        """注册工作流"""
        if name in self.workflow_registry:
            self.logger.warning(f"工作流 {name} 已存在，将被覆盖")
        
        self.workflow_registry[name] = agent_sequence
        self.logger.info(f"注册工作流: {name} -> {agent_sequence}")
    
    def execute_workflow(self, workflow_name: str, initial_input: Dict[str, Any]) -> Dict[str, Any]:
        """执行工作流"""
        if workflow_name not in self.workflow_registry:
            raise ValueError(f"工作流 '{workflow_name}' 未注册")
        
        agent_sequence = self.workflow_registry[workflow_name]
        current_input = initial_input.copy()
        execution_results = {}
        
        self.logger.info(f"开始执行工作流: {workflow_name}")
        workflow_start_time = datetime.now()
        
        for agent_name in agent_sequence:
            if agent_name not in self.agents:
                raise ValueError(f"智能体 '{agent_name}' 未注册")
            
            agent = self.agents[agent_name]
            
            # 检查依赖是否满足
            dependencies = agent.get_dependencies()
            for dep in dependencies:
                if dep not in execution_results:
                    raise ValueError(f"智能体 '{agent_name}' 依赖的 '{dep}' 未执行")
            
            self.logger.info(f"执行智能体: {agent_name}")
            
            try:
                start_time = datetime.now()
                result: AgentResult = agent.execute(current_input)
                execution_time = (datetime.now() - start_time).total_seconds()
                
                execution_results[agent_name] = {
                    "result": result,
                    "execution_time": execution_time,
                    "timestamp": datetime.now().isoformat()
                }
                
                if result.success:
                    # 将结果合并到当前输入中，供后续智能体使用
                    current_input[agent_name + "_result"] = result.data
                    current_input.update(result.metadata or {})
                    self.logger.info(f"智能体 {agent_name} 执行成功，耗时: {execution_time:.2f}s")
                else:
                    self.logger.error(f"智能体 {agent_name} 执行失败: {result.error}")
                    # 可以根据策略决定是否继续执行
                    break
                    
            except Exception as e:
                self.logger.error(f"智能体 {agent_name} 执行异常: {str(e)}")
                execution_results[agent_name] = {
                    "result": AgentResult(success=False, data=None, error=str(e)),
                    "execution_time": 0,
                    "timestamp": datetime.now().isoformat()
                }
                break
        
        workflow_time = (datetime.now() - workflow_start_time).total_seconds()
        
        # 计算总体成功率
        successful_agents = [r for r in execution_results.values() if r["result"].success]
        success_rate = len(successful_agents) / len(execution_results) if execution_results else 0
        
        final_result = {
            "workflow_name": workflow_name,
            "success": all(r["result"].success for r in execution_results.values()),
            "success_rate": success_rate,
            "execution_results": execution_results,
            "final_output": current_input,
            "workflow_execution_time": workflow_time,
            "timestamp": datetime.now().isoformat()
        }
        
        self.logger.info(f"工作流 {workflow_name} 执行完成，成功率: {success_rate:.1%}，总耗时: {workflow_time:.2f}s")
        
        return final_result
    
    def get_agent_status(self) -> Dict[str, Any]:
        """获取智能体状态"""
        status = {}
        for name, agent in self.agents.items():
            status[name] = {
                "name": name,
                "version": agent.get_version(),
                "dependencies": agent.get_dependencies()
            }
        return status
    
    def get_available_workflows(self) -> Dict[str, List[str]]:
        """获取可用工作流"""
        return self.workflow_registry.copy()
    
    def validate_workflow(self, workflow_name: str) -> bool:
        """验证工作流配置"""
        if workflow_name not in self.workflow_registry:
            return False
        
        agent_sequence = self.workflow_registry[workflow_name]
        
        # 检查所有智能体是否都已注册
        for agent_name in agent_sequence:
            if agent_name not in self.agents:
                return False
        
        return True
